#include "kpextOpImgCvtF32C1toU16C1.h"

using namespace kpext;

OpImgCvtF32C1toU16C1::OpImgCvtF32C1toU16C1(kpextVkContext *ctx) :
    BaseCompute(ctx)
{
    if(!m_vkCtx){
        throw std::runtime_error(__FUNCTION__ " construct with null pointer");
    }

    m_cols = 1;
    m_rows = 1;
    m_inBuf = nullptr;
    m_outBuf = nullptr;

    m_postBarrier = true;

    //  初始化描述符集及布局
    m_descrPool = m_vkCtx->descriptorPool(
                {2},
                1);

    vk::DescriptorSetLayoutBinding binding(
                0,
                vk::DescriptorType::eStorageBuffer,
                1,
                vk::ShaderStageFlagBits::eCompute
                );
    std::vector<vk::DescriptorSetLayoutBinding> bindings;
    bindings.push_back(
                binding
                );
    binding.setBinding(1);
    bindings.push_back(binding);

    m_descrSetlayout = m_descrPool->createDescriptorSetLayout(
                bindings
                );

    m_descrSet = m_descrPool->recreateDescSet(*m_descrSetlayout,
                                              m_descrSet
                                              );

    // 初始化pipeline
    m_pipeline = m_vkCtx->computePipeline();
    m_pipeline->m_specConsts = std::make_shared<PipelineConstSZ>(3, 4);
    m_pipeline->m_pushConsts = std::make_shared<PipelineConstSZ>(2, 4);

    m_workGroupSize = ctx->getDefWorkGroupSize();

    m_pipeline->m_specConsts->set({m_workGroupSize[0] * m_workGroupSize[1],
                                   1,
                                   1
                                  });
    m_pipeline->m_pushConsts->set({1, 1});

    m_pipeline->m_descrSetLayouts.push_back(m_descrSetlayout);

    if(!m_pipeline->rebuild(
                sizeof (spirv::ImgCvtF32C1toU16C1), spirv::ImgCvtF32C1toU16C1)
            ){
        SPDLOG_WARN("Rebuild kpext::ComputePipeline failed !");
    }

    //  工作组数量
    m_workCount = {1, 1, 1};
}

OpImgCvtF32C1toU16C1::~OpImgCvtF32C1toU16C1()
{

}

void OpImgCvtF32C1toU16C1::rebuild(int cols,
                                   int rows,
                                   std::shared_ptr<BaseBuffer> inbuf,
                                   std::shared_ptr<BaseBuffer> outbuf)
{
    if(!inbuf || !outbuf || !m_vkCtx){
        throw std::runtime_error(__FUNCTION__ " input with null pointer!");
    }

    if(cols < 0) cols = 1;
    if(rows < 0) rows = 1;

    m_cmdBarrier.prepare(
                {outbuf,},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eShaderRead
                );

    STCAST_ASSIGN(m_cols, cols);
    STCAST_ASSIGN(m_rows, rows);

    // 检查缓冲尺寸
    uint32_t t = m_cols * m_rows;
    t = (t + 1) / 2;
    if(inbuf->getSize() / 4 < static_cast<uint32_t>(t * 2)
            || outbuf->getSize() / 4 < static_cast<uint32_t>(t)
            ){
        throw std::runtime_error(__FUNCTION__ " input or output buffer size not enought!");
    }

    if(m_inBuf != inbuf.get()
            || m_outBuf != outbuf.get()
            ){
        BaseCompute::rebuild({inbuf, outbuf});
        m_inBuf = inbuf.get();
        m_outBuf = outbuf.get();
    }

    //  设定输入输出尺寸
    m_pipeline->m_pushConsts->set(
                {m_cols * m_rows, t
                }
                );

    //  重新计算工作组大小
    m_workCount[1] = 1;
    m_workCount[2] = 1;
    auto t1 = m_workGroupSize[0] * m_workGroupSize[1];
    m_workCount[0] = (t + t1 - 1) / t1;
    const uint32_t MaxWGX =
            m_vkCtx->getMaxWorkGroupCount()[0]; //4096;
    if(m_workCount[0] > MaxWGX){
        m_workCount[0] = MaxWGX;
        t1 = MaxWGX * t1;
        m_workCount[1] = (t + t1 - 1) / t1;
    }
}
